package com.gitee.yusugar.core.ssh;

import cn.hutool.core.io.IoUtil;
import cn.hutool.extra.ssh.ChannelType;
import cn.hutool.extra.ssh.JschUtil;
import com.jcraft.jsch.Channel;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import lombok.extern.slf4j.Slf4j;

import java.io.OutputStream;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

/**
 * @author Y.sugar
 * @date 2023/1/6 17:38
 */
@Slf4j
public class SSHInfo {

    private final Session session;

    private final ExecutorService singleThread = Executors.newSingleThreadScheduledExecutor();

    public SSHInfo(String addr, Integer port, String user, String pass) {
        session = JschUtil.getSession(addr, port, user, pass);
    }

    public Session getSession(){
        return this.session;
    }

    /**
     * 获取 session是否关闭
     * @return true关闭 false未关闭
     */
    public boolean sessionClosed() {
        return !session.isConnected();
    }

    /**
     * 开启 session 连接
     * @return 是否开启成功
     */
    public boolean openSession() {
        try {
            if (sessionClosed())
                session.connect();
            return true;
        } catch (JSchException e) {
            log.error("session连接失败");
            return false;
        }
    }

    /**
     * 向设备发送命令
     * @param cmdStrList 命令集合
     * @return 设备返回的执行结果
     */
    public String sendCmdStr(List<String> cmdStrList) {
        return sendCmdStr(cmdStrList.toArray());
    }


    /**
     * 向设备发送命令
     * @param cmdStrArr 命令集合
     * @return 设备返回的执行结果
     */
    public String sendCmdStr(Object... cmdStrArr) {
        AtomicReference<String> execRes = new AtomicReference<>();

        // 设置等待一个线程数，以防导致ssh连接没有关闭，流会一直处于阻塞状态的问题
        CountDownLatch countDownLatch = new CountDownLatch(1);

        AtomicReference<OutputStream> os = new AtomicReference<>();
        AtomicReference<Channel> channel = new AtomicReference<>();

        singleThread.execute(() -> {
            try {
                // 创建通道
                channel.set(JschUtil.createChannel(session, ChannelType.SHELL));
                channel.get().connect();

                // 获取输出流并写出命令
                os.set(channel.get().getOutputStream());
                IoUtil.writeUtf8(os.get(), false, cmdStrArr);

                // 获取输入流并返回结果
                execRes.set(IoUtil.readUtf8(channel.get().getInputStream()));
            } catch (Exception e) {
                log.error("Ip地址为：{} 发送失败.", session.getHost());
            } finally {
                if (os.get() != null) {
                    IoUtil.close(os.get());
                }
                if (channel.get() != null) {
                    JschUtil.close(channel.get());
                }
                countDownLatch.countDown();
            }
        });

        try {
            // 在这判断是否等待失败，等待失败后需将原先创建的流和通道关闭，不然会影响后面再次发送命令
            boolean awaitSuc = countDownLatch.await(2, TimeUnit.SECONDS);
            if (!awaitSuc) {
                if (os.get() != null) {
                    IoUtil.close(os.get());
                }
                if (channel.get() != null) {
                    JschUtil.close(channel.get());
                }
            }
        } catch (Exception e) {
            log.error("Ip地址为：{} 线程等待失败.", session.getHost());
            return null;
        }

        return execRes.get();
    }
}
